import os
import torch
import torchvision
import torchvision.transforms as transforms
from torch import nn
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
from torchvision.utils import save_image
import random


# ========================= 数据加载与预处理 =========================
def load_data(batch_size=64):
    transform = transforms.ToTensor()
    train_data = torchvision.datasets.MNIST(root='data', train=True, download=True, transform=transform)
    test_data = torchvision.datasets.MNIST(root='data', train=False, download=True, transform=transform)
    train_dataloader = DataLoader(dataset=train_data, batch_size=batch_size)
    test_dataloader = DataLoader(dataset=test_data, batch_size=batch_size)
    return train_dataloader, test_dataloader


# ========================= 模型定义 =========================
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 6, 3, stride=1, padding=1),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5, stride=1, padding=1),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(576, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


# ========================= 训练函数 =========================
def train_model(network, train_dataloader, device, loss_fn, optimizer, epochs=10):
    network.train()
    losses = []
    for epoch in range(epochs):
        loss_sum = 0
        for X, y in train_dataloader:
            X, y = X.to(device), y.to(device)
            pred = network(X)
            loss = loss_fn(pred, y)
            loss_sum += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        mean_loss = loss_sum / len(train_dataloader.dataset)
        losses.append(mean_loss)
        print(f"Epoch {epoch + 1} loss: {mean_loss:.6f}")
    torch.save(network.state_dict(), "model.pth")
    plt.xlabel("Epochs")
    plt.ylabel("Loss Value")
    plt.plot(range(epochs), losses)
    plt.savefig("training_loss.png")
    print("Training loss plot saved to training_loss.png")


# ========================= 测试函数 =========================
def evaluate_model(network, test_dataloader, device):
    network.eval()
    positive, negative = 0, 0
    with torch.no_grad():
        for X, y in test_dataloader:
            X, y = X.to(device), y.to(device)
            pred = network(X)
            for item in zip(pred, y):
                if torch.argmax(item[0]) == item[1]:
                    positive += 1
                else:
                    negative += 1
    acc = positive / (positive + negative)
    print(f"Accuracy: {acc * 100:.2f}%")
    return acc


# ========================= DeepFool攻击算法实现 =========================
def deepfool_attack_single(model, image, num_classes=10, max_iter=100, overshoot=0.1):
    """
    对单个图像执行DeepFool攻击
    
    参数:
        model: 模型
        image: 输入图像张量(已扩展批次维度)
        num_classes: 类别数
        max_iter: 最大迭代次数 (增加到100)
        overshoot: 越界系数 (增加到0.1以增强攻击)
    
    返回:
        对抗样本
        扰动
    """
    model.eval()  # 确保模型处于评估模式
    image = image.detach().clone()
    
    with torch.no_grad():
        f_image = model(image).flatten()
        I = f_image.argsort(descending=True)
        label = I[0]
    
    input_shape = image.shape
    x = image.detach().clone()
    x.requires_grad = True    
    w = torch.zeros(input_shape).to(image.device)
    r_tot = torch.zeros(input_shape).to(image.device)    
    fs = model(x)[0]
    k_i = label   
    iter_count = 0
    
    while k_i == label and iter_count < max_iter:
        pert = float('inf')
        # 清除之前的梯度
        if x.grad is not None:
            x.grad.zero_()            
        fs[label].backward(retain_graph=True)
        grad_orig = x.grad.clone()        
        for k in range(1, num_classes):
            # 清零梯度
            x.grad.zero_()            
            # 计算其他类的输出
            fs[I[k]].backward(retain_graph=True)
            grad_k = x.grad.clone()            
            # 计算梯度差
            w_k = grad_k - grad_orig
            w_k_norm = torch.norm(w_k.flatten()) + 1e-10            
            # 计算决策边界距离
            f_k = (fs[I[k]] - fs[label]).item()
            pert_k = abs(f_k) / w_k_norm            
            # 寻找最近的决策边界
            if pert_k < pert:
                pert = pert_k
                w = w_k
        
        # 计算最小扰动
        r_i = pert * w / (torch.norm(w.flatten()) + 1e-10)
        r_tot = r_tot + r_i        
        # 添加扰动
        x = image + (1 + overshoot) * r_tot
        x = torch.clamp(x, 0, 1)
        x.requires_grad = True        
        fs = model(x)[0]
        k_i = torch.argmax(fs).item()
        iter_count += 1
    
    r_tot = (1 + overshoot) * r_tot
    return image + r_tot, r_tot


# ========================= 无约束DeepFool攻击 =========================
def deepfool_attack_unconstrained(network, test_dataloader, device, output_dir="deepfool_samples"):
    """
    执行原始无约束的DeepFool攻击
    """
    network.eval()
    positive, negative = 0, 0
    perturbation_norms = []
    
    # 创建输出文件夹
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    for batch_idx, (X, y) in enumerate(test_dataloader):
        X, y = X.to(device), y.to(device)
        X_adv_batch = torch.zeros_like(X)
        
        # 对批次中的每个图像单独进行DeepFool攻击
        for i in range(X.shape[0]):
            try:
                X_adv, r = deepfool_attack_single(
                    network, X[i:i+1], num_classes=10, max_iter=100, overshoot=0.1
                )
                X_adv_batch[i] = X_adv
                perturbation_norms.append(torch.norm(r.flatten(), p=2).item())
            except:
                # 如果攻击失败，使用原始图像
                X_adv_batch[i] = X[i]
        
        # 评估对抗样本
        with torch.no_grad():
            pred = network(X_adv_batch)
            for item in zip(pred, y):
                if torch.argmax(item[0]) == item[1]:
                    positive += 1
                else:
                    negative += 1
        
        # 保存图像
        if batch_idx == 0:
            original_images_path = os.path.join(output_dir, "original_images.png")
            save_image(X.cpu(), original_images_path, nrow=8)
            print(f"Saved original images to {original_images_path}")
            
            adversarial_images_path = os.path.join(output_dir, "deepfool_unconstrained_adversarial.png")
            save_image(X_adv_batch.cpu(), adversarial_images_path, nrow=8)
            print(f"Saved unconstrained DeepFool adversarial samples")
            
            # 保存扰动可视化
            perturbation = (X_adv_batch - X).cpu()
            perturbation = perturbation * 10  # 对无约束扰动放大10倍以便观察
            perturbation = torch.clamp(perturbation + 0.5, 0, 1)
            pert_path = os.path.join(output_dir, "deepfool_unconstrained_perturbation.png")
            save_image(perturbation, pert_path, nrow=8)
            print(f"Saved unconstrained perturbation visualization")
        
        # 只处理前几个批次
        if batch_idx >= 5:
            break
    
    acc = positive / (positive + negative)
    avg_norm = np.mean(perturbation_norms) if perturbation_norms else 0
    
    print(f"无约束DeepFool攻击准确率: {acc*100:.2f}%")
    print(f"平均扰动范数: {avg_norm:.4f}")
    
    return acc, avg_norm


# ========================= 主函数 =========================
def main():

    # 加载数据
    batch_size = 64
    train_dataloader, test_dataloader = load_data(batch_size)

    # 设置设备
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")

    # 初始化模型、损失函数和优化器
    network = LeNet().to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(network.parameters(), lr=0.001, momentum=0.9)

    # 训练或加载模型
    if os.path.exists('model.pth'):
        network.load_state_dict(torch.load('model.pth'))
        print("Loaded pre-trained model.")
    else:
        print("Training model...")
        train_model(network, train_dataloader, device, loss_fn, optimizer)

    # 测试模型
    print("Evaluating model...")
    evaluate_model(network, test_dataloader, device)

    # 执行无约束DeepFool攻击
    print("执行无约束DeepFool攻击")
    deepfool_attack_unconstrained(network, test_dataloader, device)
    

if __name__ == "__main__":
    main()